from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *

import os
import xlwt
import xlrd
import xlsxwriter

class DataTabView(QTableView):
    def __init__(self, parent=None):
        '''表格初始化'''
        super(DataTabView, self).__init__(parent)
        self.model = QStandardItemModel(0, 0)
        #  self.HeaderList = ['x(mm)', 'y(mm)', 'z(mm)', '主测通道(mGy)', '监测通道(mGy)', '修正', '归一化']
        self.HeaderList = ['x(mm)', 'y(mm)', 'z(mm)', '主测通道(cGy/min)', '监测通道(cGy/min)']
        self.model.setHorizontalHeaderLabels(self.HeaderList)  #
        self.setModel(self.model)
        # 下面代码让表格100填满窗口
        self.horizontalHeader().setStretchLastSection(True)
        self.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.isClearChooseLine_Flag = False
        #  self.setMaximumHeight(250)

        self.clear_data()
        # 右键菜单设置
        self.setContextMenuPolicy(Qt.CustomContextMenu)
        self.customContextMenuRequested.connect(self.right_menu)

    def right_menu(self):
        self.menu = QMenu()
        exportAction = QAction("&导出excel", triggered=self.exportToExcel)
        self.menu.addAction(exportAction)
        self.menu.exec(QCursor.pos())

    def exportToExcel(self):
        f = QFileDialog.getSaveFileName(self, "导出Excel", '%s/data' % os.getcwd(), 'Excel Files(*.xls)')
        if f[0]:
            try:
                filename = f[0]
                wb = xlwt.Workbook()
                ws = wb.add_sheet('sheet1')  # sheetname
                data = []
                data.append(self.HeaderList)
                row = self.model.rowCount()
                for r in range(0, row):
                    x = self.model.data(self.model.index(r, 0))
                    y = self.model.data(self.model.index(r, 1))
                    z = self.model.data(self.model.index(r, 2))
                    t1 = self.model.data(self.model.index(r, 3))
                    t2 = self.model.data(self.model.index(r, 4))
                    if x and y and z and t1 and t2:
                        data = []
                        for r in range(0,row):
                            data.append([float(x), float(y), float(z), float(t1), float(t2)])
                        if data:
                            for i in range(0, len(data)):
                                for j in range(0, len(data[i])):
                                    ws.write(i, j, data[i][j])
                            wb.save(filename)
                            QMessageBox.warning(self, '提示', '导出成功!\r\n%s' % filename, QMessageBox.Yes)
                            return
                QMessageBox.warning(self, '警告', '导出失败!', QMessageBox.Yes)
            except:
                QMessageBox.warning(self, '警告', '无法导出!', QMessageBox.Yes)

    def clear_data(self):
        self.model.clear()
        self.model = QStandardItemModel(0, 0)
        self.model.setHorizontalHeaderLabels(self.HeaderList)  #
        self.setModel(self.model)

    def add_data(self, x='', y='', z='', t1='', t2=''):
        rowNum = self.model.rowCount()  # 总行数
        self.model.setItem(rowNum, 0, QStandardItem(str(x)))
        self.model.setItem(rowNum, 1, QStandardItem(str(y)))
        self.model.setItem(rowNum, 2, QStandardItem(str(z)))
        self.model.setItem(rowNum, 3, QStandardItem(str('%.2f'%t1)))
        self.model.setItem(rowNum, 4, QStandardItem(str('%.2f'%t2)))

    def insert_data(self, rowNum, x='', y='', z='', t1='', t2=''):
        # rowNum = self.model.rowCount()  # 总行数
        self.model.setItem(rowNum, 0, QStandardItem(str(x)))
        self.model.setItem(rowNum, 1, QStandardItem(str(y)))
        self.model.setItem(rowNum, 2, QStandardItem(str(z)))
        self.model.setItem(rowNum, 3, QStandardItem(str('%.2f'%t1)))
        self.model.setItem(rowNum, 4, QStandardItem(str('%.2f'%t2)))

    def get_standard(self, x, ave):
        import math
        n = len(x)
        s = 0
        for i in range(0, n):
            s = s + pow((int(x[i]) - ave), 2)
        s = math.sqrt(1 / (n - 1) * s)
        return s

    def get_data(self):
        rowNum = self.model.rowCount()  # 总行数
        x = [self.model.data(self.model.index(i, 1)) for i in range(0, rowNum)]
        line1 = [self.model.data(self.model.index(i, 4)) for i in range(0, rowNum)]
        line2 = [self.model.data(self.model.index(i, 5)) for i in range(0, rowNum)]
        return (x, line1, line2)

    def get_all_data(self):
        rowNum = self.model.rowCount()  # 总行数
        colNum = self.model.columnCount()  # 总行数
        data = []
        for j in range(0, colNum):
            col = [self.model.data(self.model.index(i, j)) for i in range(0, rowNum)]
            data.append(col)
        return data

    def get_datas(self):
        rowNum = self.model.rowCount()  # 总行数
        colNum = self.model.columnCount()  # 总行数
        data = []
        for i in range(0, rowNum):
            row_data = [float(self.model.data(self.model.index(i, j))) for j in range(0, colNum)]
            data.append(row_data)
        return data